{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "405380d381534be1b0d92a5da1013632": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_54eb133663084cde8c32854b47558f83",
              "IPY_MODEL_c74c928ca40d4215b600f617578c10b3",
              "IPY_MODEL_71021babdbe140e480172a84d4d69749"
            ],
            "layout": "IPY_MODEL_9215306114b24fffa4223f896cd50e1d"
          }
        },
        "54eb133663084cde8c32854b47558f83": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_46a0ba39225d4884aad26d72d8bd0c72",
            "placeholder": "​",
            "style": "IPY_MODEL_688a5aad1e8c4188ac5c5c705832b066",
            "value": "100%"
          }
        },
        "c74c928ca40d4215b600f617578c10b3": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_b1c5b6d6f4244c02964caa104dc0b170",
            "max": 509663,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_80e40812cfcf4a25bec8dc1d2e6a18ca",
            "value": 509663
          }
        },
        "71021babdbe140e480172a84d4d69749": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_441c174815e443698552dd5ee5a82bb3",
            "placeholder": "​",
            "style": "IPY_MODEL_70dc83747ccc4b85b4a7b90cdeb5a301",
            "value": " 509663/509663 [00:17&lt;00:00, 16838.02it/s]"
          }
        },
        "9215306114b24fffa4223f896cd50e1d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "46a0ba39225d4884aad26d72d8bd0c72": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "688a5aad1e8c4188ac5c5c705832b066": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "b1c5b6d6f4244c02964caa104dc0b170": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "80e40812cfcf4a25bec8dc1d2e6a18ca": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "441c174815e443698552dd5ee5a82bb3": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "70dc83747ccc4b85b4a7b90cdeb5a301": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "In the past months, we engineered a novel relevance endpoint that takes a query and a list of documents and predicts the relevance between the query and each document. \n",
        "\n",
        "It can be used in a two-stage retrieval setup: First you take the user question, and retrieve the top-100 documents from your collection by either using lexical search or semantic search.\n",
        "\n",
        "You then pass the question and these top-100 documents to our relevance-endpoint to get a score for each document. You can then rank these documents based on these scores.\n",
        "\n",
        "In our benchmarks across 20 datasets, we **saw significant improvements compared to lexical and semantic search**, especially for use-cases where no training data is available."
      ],
      "metadata": {
        "id": "xjLFw9Vfi3xO"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "We will demonstrate the rerank endpoint in this notebook.\n",
        "\n"
      ],
      "metadata": {
        "id": "Jw7C9tZxVgXm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# TODO: upgrade to \"cohere>5\"",
"!pip install \"cohere<5\""
      ],
      "metadata": {
        "id": "ZDfo_Y0Wjd9o"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "07xZdGGAOCaP"
      },
      "outputs": [],
      "source": [
        "import cohere\n",
        "import requests\n",
        "import numpy as np\n",
        "from time import time\n",
        "from typing import List\n",
        "from pprint import pprint"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Set up your cohere client\n",
        "API_KEY = \"<YOUR-API-KEY>\"\n",
        "co = cohere.Client(API_KEY)\n",
        "MODEL_NAME = \"rerank-english-02\" # another option is rerank-multilingual-02\n",
        "\n",
        "# Example query and passages (data taken from http://sbert.net/datasets/simplewiki-2020-11-01.jsonl.gz)\n",
        "query = \"What is the capital of the United States?\"\n",
        "docs = [\n",
        "    \"Carson City is the capital city of the American state of Nevada. At the 2010 United States Census, Carson City had a population of 55,274.\",\n",
        "    \"The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean that are a political division controlled by the United States. Its capital is Saipan.\",\n",
        "    \"Charlotte Amalie is the capital and largest city of the United States Virgin Islands. It has about 20,000 people. The city is on the island of Saint Thomas.\",\n",
        "    \"Washington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) is the capital of the United States. It is a federal district. The President of the USA and many major national government offices are in the territory. This makes it the political center of the United States of America.\",\n",
        "    \"West Virginia is a state in the Appalachian region of the United States. Its capital and largest city is Charleston. It is often abbreviated W. Va. or simply WV.\",\n",
        "    \"Capital punishment (the death penalty) has existed in the United States since before the United States was a country. As of 2017, capital punishment is legal in 30 of the 50 states. The federal government (including the United States military) also uses capital punishment.\",\n",
        "    \"North Dakota is a state in the United States. 672,591 people lived in North Dakota in the year 2010. The capital and seat of government is Bismarck.\",\n",
        "    \"Kentucky is a state in the United States. Its capital is Frankfort. It touches the states of Missouri (by the Mississippi River), Illinois, Indiana, Ohio, West Virginia (by the Ohio River), Tennessee and Virginia. There are many rivers in Kentucky\",\n",
        "    \"Micronesia, officially the Federated States of Micronesia, is an island nation in the Pacific Ocean, northeast of Papua New Guinea. The country is a sovereign state in free association with the United States. The capital city of Federated States of Micronesia is Palikir.\",\n",
        "    \"Utah is a state in the west United States. The capital and largest city is Salt Lake City. Utah became a state in the U.S. on January 4, 1896.\"]"
      ],
      "metadata": {
        "id": "FqKP7G3fOVQ9"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Using the Endpoint\n",
        "In the following cell we will call rerank to rank `docs` based on how relevant they are with `query`.\n",
        "\n"
      ],
      "metadata": {
        "id": "u3Ogs_QotGTw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "results = co.rerank(query=query, model=MODEL_NAME, documents=docs, top_n=3) # Change top_n to change the number of results returned. If top_n is not passed, all results will be returned.\n",
        "for idx, r in enumerate(results):\n",
        "  print(f\"Document Rank: {idx + 1}, Document Index: {r.index}\")\n",
        "  print(f\"Document: {r.document['text']}\")\n",
        "  print(f\"Relevance Score: {r.relevance_score:.2f}\")\n",
        "  print(\"\\n\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "05S-LK5uthHz",
        "outputId": "7ca9af1c-b682-4d8f-d062-7dcdb267dbf4"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Document Rank: 1, Document Index: 3\n",
            "Document: Washington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) is the capital of the United States. It is a federal district. The President of the USA and many major national government offices are in the territory. This makes it the political center of the United States of America.\n",
            "Relevance Score: 0.99\n",
            "\n",
            "\n",
            "Document Rank: 2, Document Index: 1\n",
            "Document: The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean that are a political division controlled by the United States. Its capital is Saipan.\n",
            "Relevance Score: 0.30\n",
            "\n",
            "\n",
            "Document Rank: 3, Document Index: 5\n",
            "Document: Capital punishment (the death penalty) has existed in the United States since before the United States was a country. As of 2017, capital punishment is legal in 30 of the 50 states. The federal government (including the United States military) also uses capital punishment.\n",
            "Relevance Score: 0.27\n",
            "\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Search on  Wikipedia - End2end demo\n",
        "The following is an example how to use this model end-to-end to search over the Simple English Wikipedia, which consists of about 500k passages. \n",
        "\n",
        "We use BM25 lexical search to retrieve the top-100 passages matching the query and then send these 100 passages and the query to our rerank endpoint to get a re-ranked list. We output the top-3 hits according to BM25 lexical search (as used by e.g. Elasticsearch) and the re-ranked list from our endpoint.\n"
      ],
      "metadata": {
        "id": "61WyvMAtlf2Y"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -U  rank_bm25"
      ],
      "metadata": {
        "id": "o-gpxa9NlkK7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "import gzip\n",
        "import os\n",
        "from rank_bm25 import BM25Okapi\n",
        "from sklearn.feature_extraction import _stop_words\n",
        "import string\n",
        "from tqdm.autonotebook import tqdm"
      ],
      "metadata": {
        "id": "YOsh4bfvlkRQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!wget http://sbert.net/datasets/simplewiki-2020-11-01.jsonl.gz"
      ],
      "metadata": {
        "id": "_Rf_d54tlvtM",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "84583161-c3cd-4f14-89d8-995d2aa55eef"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2023-04-17 22:58:37--  http://sbert.net/datasets/simplewiki-2020-11-01.jsonl.gz\n",
            "Resolving sbert.net (sbert.net)... 172.67.180.145, 104.21.67.200, 2606:4700:3036::6815:43c8, ...\n",
            "Connecting to sbert.net (sbert.net)|172.67.180.145|:80... connected.\n",
            "HTTP request sent, awaiting response... 301 Moved Permanently\n",
            "Location: https://sbert.net/datasets/simplewiki-2020-11-01.jsonl.gz [following]\n",
            "--2023-04-17 22:58:38--  https://sbert.net/datasets/simplewiki-2020-11-01.jsonl.gz\n",
            "Connecting to sbert.net (sbert.net)|172.67.180.145|:443... connected.\n",
            "HTTP request sent, awaiting response... 301 Moved Permanently\n",
            "Location: https://public.ukp.informatik.tu-darmstadt.de/reimers/sentence-transformers/datasets/simplewiki-2020-11-01.jsonl.gz [following]\n",
            "--2023-04-17 22:58:38--  https://public.ukp.informatik.tu-darmstadt.de/reimers/sentence-transformers/datasets/simplewiki-2020-11-01.jsonl.gz\n",
            "Resolving public.ukp.informatik.tu-darmstadt.de (public.ukp.informatik.tu-darmstadt.de)... 130.83.167.186\n",
            "Connecting to public.ukp.informatik.tu-darmstadt.de (public.ukp.informatik.tu-darmstadt.de)|130.83.167.186|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 50223724 (48M) [application/octet-stream]\n",
            "Saving to: ‘simplewiki-2020-11-01.jsonl.gz’\n",
            "\n",
            "simplewiki-2020-11- 100%[===================>]  47.90M  20.1MB/s    in 2.4s    \n",
            "\n",
            "2023-04-17 22:58:41 (20.1 MB/s) - ‘simplewiki-2020-11-01.jsonl.gz’ saved [50223724/50223724]\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# As dataset, we use Simple English Wikipedia. Compared to the full English wikipedia, it has only\n",
        "# about 170k articles. We split these articles into paragraphs and encode them with the bi-encoder\n",
        "wikipedia_filepath = 'simplewiki-2020-11-01.jsonl.gz'\n",
        "\n",
        "passages = []\n",
        "with gzip.open(wikipedia_filepath, 'rt', encoding='utf8') as fIn:\n",
        "    for line in fIn:\n",
        "        data = json.loads(line.strip())\n",
        "        passages.extend(data['paragraphs'])\n",
        "\n",
        "print(\"Passages:\", len(passages))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OMLzTvk7lzyY",
        "outputId": "1d05aa18-ec4c-4b64-d21d-76faccf4d951"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Passages: 509663\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(passages[0], passages[1])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CAtpe51XgZ5Y",
        "outputId": "1ccf0358-6712-4c21-e831-c3b3611597fd"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Ted Cassidy (July 31, 1932 - January 16, 1979) was an American actor. He was best known for his roles as Lurch and Thing on \"The Addams Family\". Aileen Carol Wuornos Pralle (born Aileen Carol Pittman; February 29, 1956 – October 9, 2002) was an American serial killer. She was born in Rochester, Michigan. She confessed to killing six men in Florida and was executed in Florida State Prison by lethal injection for the murders. Wuornos said that the men she killed had raped her or tried to rape her while she was working as a prostitute.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# We compare the results to lexical search (keyword search). Here, we use \n",
        "# the BM25 algorithm which is implemented in the rank_bm25 package.\n",
        "\n",
        "# We lower case our text and remove stop-words from indexing\n",
        "def bm25_tokenizer(text):\n",
        "    tokenized_doc = []\n",
        "    for token in text.lower().split():\n",
        "        token = token.strip(string.punctuation)\n",
        "\n",
        "        if len(token) > 0 and token not in _stop_words.ENGLISH_STOP_WORDS:\n",
        "            tokenized_doc.append(token)\n",
        "    return tokenized_doc\n",
        "\n",
        "\n",
        "tokenized_corpus = []\n",
        "for passage in tqdm(passages):\n",
        "    tokenized_corpus.append(bm25_tokenizer(passage))\n",
        "\n",
        "bm25 = BM25Okapi(tokenized_corpus)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "405380d381534be1b0d92a5da1013632",
            "54eb133663084cde8c32854b47558f83",
            "c74c928ca40d4215b600f617578c10b3",
            "71021babdbe140e480172a84d4d69749",
            "9215306114b24fffa4223f896cd50e1d",
            "46a0ba39225d4884aad26d72d8bd0c72",
            "688a5aad1e8c4188ac5c5c705832b066",
            "b1c5b6d6f4244c02964caa104dc0b170",
            "80e40812cfcf4a25bec8dc1d2e6a18ca",
            "441c174815e443698552dd5ee5a82bb3",
            "70dc83747ccc4b85b4a7b90cdeb5a301"
          ]
        },
        "id": "8AJFsclpl8wx",
        "outputId": "4fbbbf2f-624c-4f88-bfc0-6cd7d31ec0cc"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/509663 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "405380d381534be1b0d92a5da1013632"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# This function will search all wikipedia articles for passages that\n",
        "# answer the query. We then re-rank using our rerank endpoint\n",
        "\n",
        "def search(query, top_k=3, num_candidates=100):\n",
        "    print(\"Input question:\", query)\n",
        "\n",
        "    ##### BM25 search (lexical search) #####\n",
        "    bm25_scores = bm25.get_scores(bm25_tokenizer(query))\n",
        "    top_n = np.argpartition(bm25_scores, -num_candidates)[-num_candidates:]\n",
        "    bm25_hits = [{'corpus_id': idx, 'score': bm25_scores[idx]} for idx in top_n]\n",
        "    bm25_hits = sorted(bm25_hits, key=lambda x: x['score'], reverse=True)\n",
        "    \n",
        "    print(f\"Top-3 lexical search (BM25) hits\")\n",
        "    for hit in bm25_hits[0:top_k]:\n",
        "        print(\"\\t{:.3f}\\t{}\".format(hit['score'], passages[hit['corpus_id']].replace(\"\\n\", \" \")))\n",
        "\n",
        "   \n",
        "    #Add re-ranking\n",
        "    docs = [passages[hit['corpus_id']] for hit in bm25_hits]\n",
        "    \n",
        "    print(f\"\\nTop-3 hits by rank-API ({len(bm25_hits)} BM25 hits re-ranked)\")\n",
        "    results = co.rerank(query=query, model=MODEL_NAME, documents=docs, top_n=top_k)\n",
        "    for hit in results:\n",
        "        print(\"\\t{:.3f}\\t{}\".format(hit.relevance_score, hit.document[\"text\"].replace(\"\\n\", \" \")))"
      ],
      "metadata": {
        "id": "sZavW8t6mNTH"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "search(query = \"What is the capital of the United States?\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LdW9KxlUmHFN",
        "outputId": "3ce44a8d-693d-493e-b477-166a7faca4db"
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Input question: What is the capital of the United States?\n",
            "Top-3 lexical search (BM25) hits\n",
            "\t16.264\tCapital punishment (the death penalty) has existed in the United States since before the United States was a country. As of 2017, capital punishment is legal in 30 of the 50 states. The federal government (including the United States military) also uses capital punishment.\n",
            "\t15.124\tIn 1783, it was the capital of the United States for a few months.\n",
            "\t14.476\tNew York was the capital of the United States under the Articles of Confederation from 1785 to 1788. When the US Constitution was made, it stayed as the capital from 1789 until 1790. In 1789, the first President of the United States, George Washington, was inaugurated; the first United States Congress and the Supreme Court of the United States each met for the first time, and the United States Bill of Rights was written, all at Federal Hall on Wall Street. By 1790, New York grew bigger than Philadelphia, so it become the biggest city in the United States. By the end of 1790, because of the Residence Act, Philadelphia became the new capital.\n",
            "\n",
            "Top-3 hits by rank-API (100 BM25 hits re-ranked)\n",
            "\t0.987\tWashington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) is the capital of the United States. It is a federal district. The President of the USA and many major national government offices are in the territory. This makes it the political center of the United States of America.\n",
            "\t0.879\tNew York was the capital of the United States under the Articles of Confederation from 1785 to 1788. When the US Constitution was made, it stayed as the capital from 1789 until 1790. In 1789, the first President of the United States, George Washington, was inaugurated; the first United States Congress and the Supreme Court of the United States each met for the first time, and the United States Bill of Rights was written, all at Federal Hall on Wall Street. By 1790, New York grew bigger than Philadelphia, so it become the biggest city in the United States. By the end of 1790, because of the Residence Act, Philadelphia became the new capital.\n",
            "\t0.835\tAs the national capital of the United States, Washington, D.C. has numerous media outlets in various mediums. Some of these media are known throughout the United States, including \"The Washington Post\" and various broadcasting networks headquartered in D.C.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "search(query = \"Number countries Europe\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cVw7ZYI0mw3r",
        "outputId": "e959d24c-a24a-49cd-f6aa-a129a5ee8b1d"
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Input question: Number countries Europe\n",
            "Top-3 lexical search (BM25) hits\n",
            "\t16.963\tECoHR' has a number of judges. The number of judges is seven normally but at the case of dealing a great issue, the number will be 21 and the judges are equally from member countries of the Council of Europe. At present, there are forty seven member countries of the Council of Europe. Each country may have one judge in the ECoHR. But, judges work independently for the ECoHR, and not for their country.\n",
            "\t14.560\tMost countries in Europe, and a few countries in Asia, have made some or all synthetic cannabinoids illegal.\n",
            "\t14.165\tMany of these countries were members of the Western European Union. Many, such as Norway, are also in Northern Europe or in Central Europe or Southern Europe.\n",
            "\n",
            "Top-3 hits by rank-API (100 BM25 hits re-ranked)\n",
            "\t0.995\tThere are at least 43 countries in Europe (the European identities of 5 transcontinental countries:Cyprus, Georgia, Kazakhstan, Russia and Turkey are disputed). Most of these countries are members of the European Union.\n",
            "\t0.967\tWithin these regions, there are up to 48 independent European countries (with the identities of 5 transcontinental countries being disputed). The largest is the Russian Federation, which covers 39% of Europe.\n",
            "\t0.958\tEurope, the planet's 6th largest continent, includes 47 countries and assorted dependencies, islands and territories.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "search(query = \"Elon Musk year birth\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "52p0BMsomxV2",
        "outputId": "cade17e7-5bd8-44ed-d18d-2f986663ef5a"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Input question: Elon Musk year birth\n",
            "Top-3 lexical search (BM25) hits\n",
            "\t22.568\tTesla, Inc. is a company based in Palo Alto, California which makes electric cars. It was started in 2003 by Martin Eberhard, Dylan Stott, and Elon Musk (who also co-founded PayPal and SpaceX and is the CEO of SpaceX). Eberhard no longer works there. Today, Elon Musk is the Chief Executive Officer (CEO). It started selling its first car, the Roadster in 2008.\n",
            "\t20.492\tElon Musk complained via Twitter about Los Angeles traffic and the same day, December 17, 2016, founded the company. It built a short test tunnel in Los Angeles.\n",
            "\t20.448\tAt the end of 2016, Musk founded The Boring Company which focuses on tunnelling and infrastructure. He mentioned Los Angeles traffic as the reason for starting this company. In March 2017 Elon Musk announced he has started another company which aims to merge human brains and computers, it is called Neuralink.\n",
            "\n",
            "Top-3 hits by rank-API (100 BM25 hits re-ranked)\n",
            "\t0.995\tElon Reeve Musk (born June 28, 1971) is a businessman and philanthropist. He was born in South Africa. He moved to Canada and later became an American citizen. Musk is the current CEO & Chief Product Architect of Tesla Motors, a company that makes electric vehicles. He is also the CEO of Solar City, a company that makes solar panels, and the CEO & CTO of SpaceX, an aerospace company. In August 2020, Bloomberg ranked Musk third among the richest people on the planet with net worth to be $115.4 billion.\n",
            "\t0.933\tIn early 2002, Musk was seeking workers for his new space company, soon to be named SpaceX. Musk found a rocket engineer Tom Mueller (later SpaceX's CTO of Propulsion). He agreed to work for Musk. That was how SpaceX was born. The first headquarters of SpaceX was in a warehouse in El Segundo, California. The company has grown rapidly since it was founded in 2002, growing from 160 workers in November 2005 to 1,100 in 2010, 3,800 workers and contractors by October 2013, nearly 5,000 by late 2015, and about 6,000 in April 2017.\n",
            "\t0.901\tElon Musk and his brother started Zip2, a software company, in 1995. In 1999 he sold it and became a millionaire. He then started X.com, which merged with the company to make PayPal. X.com was then renamed to PayPal, and he focused on growing that part of the company. He then started SpaceX and became the CEO of Tesla.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "search(query = \"Which US president was killed?\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "r38_cEjgmzmm",
        "outputId": "1862a573-b7b0-4a3e-d720-c891fee45ea3"
      },
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Input question: Which US president was killed?\n",
            "Top-3 lexical search (BM25) hits\n",
            "\t11.966\tHe came into office when the previous president, Cyprien Ntaryamira, was killed in a plane crash. It was an assassination in which the Rwandan president Juvénal Habyarimana was also killed. Ntibantunganya left office when he was deposed by Pierre Buyoya in a military coup of 1996.\n",
            "\t11.697\tBurr killed Alexander Hamilton in a duel in 1804, when Burr was still Vice President.\n",
            "\t11.482\tAfter President James A. Garfield died, vice-president Chester Arthur replaced him. The man who killed him expected the new President to pardon him. This did not happen.\n",
            "\n",
            "Top-3 hits by rank-API (100 BM25 hits re-ranked)\n",
            "\t0.942\tPresident William McKinley was killed by anarchist Leon Czolgosz because Czolgosz believed president McKinley was against good working people, he considered McKinley responsible for falsifying the reasons for the war, and approving and waging an illegal, devastating Philippines war.\n",
            "\t0.833\tJames Abram Garfield (November 19, 1831 - September 19, 1881) was the 20th (1881) President of the United States and the 2nd President to be assassinated (killed while in office). President Garfield was in office from March to September of 1881. He was in office for a total of six months and fifteen days. For almost half that time he was bedridden as a result of an attempt to kill him. He was shot on July 2 and finally died in September the same year he got into office.\n",
            "\t0.770\tReagan was nearly killed in an assassination attempt that happened on Monday, March 30, 1981. 69 days after becoming President, he was leaving after a speaking engagement at the Washington Hilton Hotel in Washington, D.C. He was shot by John Hinckley. Hinckley shot six bullets.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "search(query=\"When is Chinese New Year\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yKc8uwkGnPIu",
        "outputId": "92aaeac2-e3ef-4621-edb1-4e1b136ed135"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Input question: When is Chinese New Year\n",
            "Top-3 lexical search (BM25) hits\n",
            "\t18.606\tToday in China the Gregorian calendar is used for most activities. At the same time, the Chinese calendar is still used for traditional Chinese holidays like Chinese New Year or Lunar New Year.\n",
            "\t18.151\tBefore that, the holiday was usually just called the \"NewYear\". Because the traditional Chinese calendar is mostly based on the changes in the moon, the Chinese New Year is also known in English as the \"Lunar New Year\" or \"Chinese Lunar New Year\". This name comes from \"Luna\", an old Latin name for the moon. The Indonesian name for the holiday is Imlek, which comes from the Hokkien word for the old Chinese calendar and is therefore also like saying \"Lunar New Year\".\n",
            "\t18.011\tSpring Festival is the Chinese New Year.\n",
            "\n",
            "Top-3 hits by rank-API (100 BM25 hits re-ranked)\n",
            "\t0.999\tChinese New Year, known in China as the SpringFestival and in Singapore as the LunarNewYear, is a holiday on and around the new moon on the first day of the year in the traditional Chinese calendar. This calendar is based on the changes in the moon and is only sometimes changed to fit the seasons of the year based on how the Earth moves around the sun. Because of this, Chinese New Year is never on January1. It moves around between January21 and February20.\n",
            "\t0.996\tChinese New Year always starts on a new moon, when the Moon is between the Earth and Sun and it looks all dark in the night sky. Because new moons happen about every 29.53 days but the year set by Pope GregoryXIII is 365.2425 days long, the Chinese holiday moves to different days each year. The Chinese calendar adds a 13th month every so often to keep the seasons in the right place, so the first day of the new year always happens between January21 and February20 on the 2nd or 3rd new moon after the 1st day of winter. The chart on the right gives the day of each Chinese New Year from 1996 to 2031.\n",
            "\t0.992\tChinese New Year lasts fifteen days, including one week as a national holiday. It starts with the first day of the Chinese lunar year and ends with the full moon fifteen days later. It is always in the middle of winter, but is called the Spring Festival in Chinese because Chinese seasons are a little different from English ones. On the first day of the Chinese New Year, people call on friends and relatives. Because most people watch the special performances on CCTV all the night on New Year's Eve and don't go to bed until 12:00 AM, they usually get up later in the next day. The fifth day of the Chinese New Year is the day to welcome the god of Wealth (Chinese:财神爷), many people make and eat dumplings (Chinese:饺子. Pinyin: Jaozi). They believe that dumplings can hold the god of Wealth and bring luck. The last day of the Chinese New Year is the Lantern Festival. On this day, the moon becomes the full moon. People go out and watch the lantern festivals everywhere. After that, they eat sweet dumpling (Chinese:汤圆,元宵), a kind of dumpling which is round and looks like the full moon.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "search(query=\"How many people live in Paris\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NT8_OxwInS1x",
        "outputId": "5ac4aba7-f26b-4a84-9b7d-41a8ab733f02"
      },
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Input question: How many people live in Paris\n",
            "Top-3 lexical search (BM25) hits\n",
            "\t16.277\tLive à Paris (English: \"Live in Paris\") is a live album by Canadian singer Céline Dion.\n",
            "\t15.173\tÎle-de-France is a region of France. The capital city is Paris. It is also the capital city of France. In 2013 about 12 million people lived in the region. About 2.1 million people live in the city of Paris.\n",
            "\t14.666\tGennevilliers is a town in France near Paris. It is in the region Île-de-France and the department of Hauts-de-Seine. About 41,000 people live there.\n",
            "\n",
            "Top-3 hits by rank-API (100 BM25 hits re-ranked)\n",
            "\t0.998\tParis (nicknamed the \"\"City of light\"\") is the capital city of France, and the largest city in France. The area is , and around 2.15 million people live there. If suburbs are counted, the population of the Paris area rises to 12 million people.\n",
            "\t0.990\tÎle-de-France is a region of France. The capital city is Paris. It is also the capital city of France. In 2013 about 12 million people lived in the region. About 2.1 million people live in the city of Paris.\n",
            "\t0.873\tMost of the 2000 students at École Centrale de Paris live in dedicated housing buildings near to the research labs and metro public transport.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "search(query=\"Who is the director of The Matrix?\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bLkBz68fnUoS",
        "outputId": "75800b8b-fe7a-4dcc-e622-b7ee7a19397a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Input question: Who is the director of The Matrix?\n",
            "Top-3 lexical search (BM25) hits\n",
            "\t16.253\tAn inverse matrix is a matrix that, when multiplied by another matrix, equals the identity matrix. For example:\n",
            "\t16.072\tis an identity matrix. There is exactly one identity matrix for each square dimension set. An identity matrix is special because when multiplying any matrix by the identity matrix, the result is always the original matrix with no change.\n",
            "\t15.353\tFirst, the system needs to be turned into an augmented matrix. In an augmented matrix, each linear equation becomes a row. On one side of the augmented matrix, the coefficients of each term in the linear equation become numbers in the matrix. On the other side of the augmented matrix are the constant terms each linear equation is equal to. For this system, the augmented matrix is:\n",
            "\n",
            "Top-3 hits by rank-API (100 BM25 hits re-ranked)\n",
            "\t5.141\tThe Matrix is a science fiction action movie that was made in 1999. It was written and directed by the Wachowski Brothers. The main actors in the movie are Keanu Reeves, Laurence Fishburne, Carrie-Anne Moss, and Hugo Weaving. \"The Matrix\" was followed by two sequels: \"The Matrix Reloaded\" and \"The Matrix Revolutions\".\n",
            "\t4.703\tHelmut Bakaitis (born 26 September 1944) is a German-born Australian director, actor and screenwriter. He is known for his role as The Architect in \"The Matrix\" movie series. Bakaitis was born in Lauban, Lower Silesia, Germany (now Lubań, Poland). Bakaitis started teaching directing at Australian Academy of Dramatic Art (AADA).\n",
            "\t3.717\tThe Matrix: Path of Neo is a 2005 action-adventure video game based on The Matrix and its sequels. It was written and directed by The Wachowskis.\n"
          ]
        }
      ]
    }
  ]
}